package com.shujia.mllib

import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.sql.functions.{count, sum, when}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo11BayesTrain {
  def main(args: Array[String]): Unit = {
    // 构建SparkSession
    val spark: SparkSession = SparkSession
      .builder()
      .appName("Demo11BayesTrain")
      .master("local[*]")
      .config("spark.sql.shuffle.partitions", "8")
      .getOrCreate()

    // 1、加载做好了数据特征工程的数据
    val rescaledData: DataFrame = spark
      .read
      .format("parquet")
      .load("Spark/data/mllib/data/rescaledData")

    // 2、将数据切分为测试集、训练集
    val arr: Array[Dataset[Row]] = rescaledData.randomSplit(Array(0.8, 0.2))
    val trainDF: Dataset[Row] = arr(0)
    val testDF: Dataset[Row] = arr(1)

    // 3、选择模型 --> 朴素贝叶斯模型
    val bayes: NaiveBayes = new NaiveBayes()

    // 4、将训练集带入模型
    val bayesModel: NaiveBayesModel = bayes.fit(trainDF)

    // 5、使用测试集评估模型
    val testResDF: DataFrame = bayesModel.transform(testDF)

    testResDF.show()

    import spark.implicits._
    // 计算模型的准确率：预测正确的/总的数据条数
    testResDF
      .withColumn("isEqual", when($"label" === $"prediction", 1).otherwise(0))
      .select(sum($"isEqual") / count("*") as "准确率")
      .show()

    // 6、保存模型
    bayesModel.save("Spark/data/mllib/bayes")


  }

}
